import gymnasium as gym
import numpy as np

np.set_printoptions(suppress=True)
#设置渲染模式为human会导致程序运行极慢
#env=gym.make("CliffWalking-v0",render_mode="human")
env = gym.make("CliffWalking-v0")

#obs=env.reset(seed=0)
#obs=obs[0]
n_state=env.observation_space.n
n_action=env.action_space.n
q_table=np.zeros((n_state,n_action),dtype=np.float32) #Q table stores Q(s,a) for each s and a.

p=0.01 #1-e greedy policy
lr=0.5 #learning  
gamma=0.99 # gamma

for i in range(200):
    num=0
    obs=env.reset(seed=0)
    obs=obs[0]
    while 1:
        state=int(obs)
        # choose an action based on 1-e greedy policy
        if np.random.uniform(0,1)>p:
            action=np.argmax(q_table[state])
        else:
            action=env.action_space.sample()
        # take a step and record data
        obs_new, r, terminated, truncated, info = env.step(action)
        state_new=int(obs_new)
        if terminated or truncated:
            break
        # update q_table using Q(s,a)<-Q(s,a)+alpha*(r+gamma*maxQ(s',a)-Q(s,a))
        q_table[state][action]=(1-lr)*q_table[state][action]+lr*(r+gamma*np.max(q_table[state_new]))
        obs=obs_new
        num+=1
    i+=1
    print("One process finished, this epoch ",i," uses :",num," steps.")

obs2=env.reset(seed=11712)
obs2=obs2[0]
num_test=0
while 1:
    state_test=int(obs2)
     # choose an action based on 1-e greedy policy
    if np.random.uniform(0,1)>p:
        a=np.argmax(q_table[state_test])
    else:
        a=env.action_space.sample()
    # take a step and record data
    obs_n, r2, ter, tru, inf = env.step(a)
    if ter or tru:
        print("Test process finished, uses :",num_test," steps.")
        break
    obs2=obs_n
    num_test+=1


'''
for _ in range(5):
    action = env.action_space.sample()  # agent policy that uses the observation and info
    print(action,type(action))
    obs, r, terminated, truncated, info = env.step(action)
    #print(obs,type(obs))
    if terminated or truncated:
        observation, info = env.reset()
        print("info is :",info)

'''

#print(q_table)
print("All done.")        
env.close()